# backend/core/wireguard.py
import subprocess
import secrets
import ipaddress
from typing import Dict, List, Optional
import logging

logger = logging.getLogger(__name__)


class WireGuardManager:
    def __init__(self, config_path: str = "/etc/wireguard"):
        self.config_path = config_path

    def generate_keypair(self) -> Dict[str, str]:
        """生成WireGuard密钥对"""
        try:
            # 生成私钥
            private_key = subprocess.run(
                ["wg", "genkey"],
                capture_output=True,
                text=True,
                check=True
            ).stdout.strip()

            # 从私钥生成公钥
            public_key = subprocess.run(
                ["wg", "pubkey"],
                input=private_key,
                capture_output=True,
                text=True,
                check=True
            ).stdout.strip()

            logger.info("✅ WireGuard密钥对生成成功")
            return {
                "private_key": private_key,
                "public_key": public_key
            }

        except subprocess.CalledProcessError as e:
            logger.error(f"❌ 生成WireGuard密钥对失败: {e}")
            # 开发环境备用方案
            return self._generate_fallback_keypair()
        except FileNotFoundError:
            logger.warning("⚠️ wg命令未找到，使用备用密钥生成方案")
            return self._generate_fallback_keypair()

    def _generate_fallback_keypair(self) -> Dict[str, str]:
        """开发环境备用密钥生成方案"""
        # 生成随机字符串作为密钥（仅用于开发）
        private_key = secrets.token_hex(32)
        public_key = secrets.token_hex(32)
        logger.info("🔧 使用备用密钥生成方案")
        return {
            "private_key": private_key,
            "public_key": public_key
        }

    def generate_client_config(
            self,
            client_private_key: str,
            server_public_key: str,
            server_endpoint: str,
            client_ip: str,
            dns_servers: List[str] = None
    ) -> str:
        """生成客户端配置文件"""
        if dns_servers is None:
            dns_servers = ["8.8.8.8", "1.1.1.1"]

        config = f"""[Interface]
PrivateKey = {client_private_key}
Address = {client_ip}
DNS = {', '.join(dns_servers)}

[Peer]
PublicKey = {server_public_key}
Endpoint = {server_endpoint}
AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 25
"""
        return config

    def allocate_client_ip(self, subnet: str, used_ips: List[str]) -> str:
        """为客户端分配IP地址"""
        network = ipaddress.ip_network(subnet)

        # 从子网中分配IP（跳过网络地址、广播地址和网关）
        for i in range(2, 254):  # 从 .2 开始分配
            ip = str(network[i])
            if ip not in used_ips:
                return ip

        raise ValueError("❌ 没有可用的IP地址")

    def validate_config(self, config: Dict) -> bool:
        """验证WireGuard配置"""
        required_fields = ['server_public_key', 'server_endpoint', 'subnet']
        return all(field in config for field in required_fields)